package org.jetbrains.kotlinx.dataframe.io

import kotlinx.datetime.LocalDate
import org.jetbrains.kotlinx.dataframe.DataColumn
import org.jetbrains.kotlinx.dataframe.api.dataFrameOf
import java.net.URI

/**
 * DataFrame to be saved in Apache Arrow
 */
val citiesExampleFrame = dataFrameOf(
    DataColumn.createValueColumn(
        "name",
        listOf(
            "Berlin",
            "Hamburg",
            "New York",
            "Washington",
            "Saint Petersburg",
            "Vatican",
        ),
    ),
    DataColumn.createValueColumn(
        "affiliation",
        listOf(
            "Germany",
            "Germany",
            "The USA",
            "The USA",
            "Russia",
            null,
        ),
    ),
    DataColumn.createValueColumn(
        "is_capital",
        listOf(
            true,
            false,
            false,
            true,
            false,
            null,
        ),
    ),
    DataColumn.createValueColumn(
        "population",
        listOf(
            3_769_495,
            1_845_229,
            8_467_513,
            689_545,
            5_377_503,
            825,
        ),
    ),
    DataColumn.createValueColumn(
        "area",
        listOf(
            891.7,
            755.22,
            1223.59,
            177.0,
            1439.0,
            0.44,
        ),
    ),
    DataColumn.createValueColumn(
        "settled",
        listOf(
            LocalDate(1237, 1, 1),
            LocalDate(1189, 5, 7),
            LocalDate(1624, 1, 1),
            LocalDate(1790, 7, 16),
            LocalDate(1703, 5, 27),
            LocalDate(1929, 2, 11),
        ),
    ),
    DataColumn.createValueColumn(
        "page_in_wiki",
        listOf(
            URI("https://en.wikipedia.org/wiki/Berlin").toURL(),
            URI("https://en.wikipedia.org/wiki/Hamburg").toURL(),
            URI("https://en.wikipedia.org/wiki/New_York_City").toURL(),
            URI("https://en.wikipedia.org/wiki/Washington,_D.C.").toURL(),
            URI("https://en.wikipedia.org/wiki/Saint_Petersburg").toURL(),
            URI("https://en.wikipedia.org/wiki/Vatican_City").toURL(),
        ),
    ),
)

/**
 * [citiesExampleFrame] Apache Arrow schema with some changes.
 * Originally generated by `citiesExampleFrame.columns().toArrowSchema().toJson()`
 * Changes made to test converting and schema matching:
 * field "population" changed to nullable Long;
 * field "area" changed to single Float;
 * field "settled" changed to datetime (date with millisecond precision);
 * field "page_in_wiki" removed, nullable field "film_in_youtube" added.
 */
val citiesExampleSchema =
    """
    {
      "fields" : [ {
        "name" : "name",
        "nullable" : false,
        "type" : {
          "name" : "utf8"
        },
        "children" : [ ]
      }, {
        "name" : "affiliation",
        "nullable" : true,
        "type" : {
          "name" : "utf8"
        },
        "children" : [ ]
      }, {
        "name" : "is_capital",
        "nullable" : true,
        "type" : {
          "name" : "bool"
        },
        "children" : [ ]
      }, {
        "name" : "population",
        "nullable" : true,
        "type" : {
          "name" : "int",
          "bitWidth" : 64,
          "isSigned" : true
        },
        "children" : [ ]
      }, {
        "name" : "area",
        "nullable" : false,
        "type" : {
          "name" : "floatingpoint",
          "precision" : "SINGLE"
        },
        "children" : [ ]
      }, {
        "name" : "settled",
        "nullable" : false,
        "type" : {
          "name" : "date",
          "unit" : "MILLISECOND"
        },
        "children" : [ ]
      }, {
        "name" : "film_in_youtube",
        "nullable" : true,
        "type" : {
          "name" : "utf8"
        },
        "children" : [ ]
      } ]
    }
    """.trimIndent()

/**
 * String column (variable length vector) with size >1 MiB
 */
val bigStringColumn = run {
    val list = ArrayList<String>()
    for (i in 0 until 1024) {
        val row = StringBuilder()
        for (j in 0 until 64) {
            row.append("abcd")
        }
        list.add(row.toString())
    }
    for (i in 0 until 1024) {
        val row = StringBuilder()
        for (j in 0 until 64) {
            row.append("гдёж")
        }
        list.add(row.toString())
    }
    for (i in 0 until 1024) {
        val row = StringBuilder()
        for (j in 0 until 64) {
            row.append("αβγδ")
        }
        list.add(row.toString())
    }
    for (i in 0 until 1024) {
        val row = StringBuilder()
        for (j in 0 until 64) {
            row.append("正体字")
        }
        list.add(row.toString())
    }
    DataColumn.createValueColumn("bigStringColumn", list)
}

val bigMixedColumn = run {
    val list = ArrayList<Any>()
    for (i in 0..32768) {
        list.add(i * i)
    }
    list.add("Dirty data")
    for (i in 32768 downTo 0) {
        list.add(i * i)
    }
    DataColumn.createValueColumn("bigMixedColumn", list)
}
